from IPython import get_ipython
from keras import Input
import keras
from keras.optimizers import SGD
from keras.utils import plot_model

from configuration.model_congig import ModelConfig
from datamanager import data_manager
from config import ModelLayers, Config
from model import Model
from log import log
from keras.applications import ResNet50V2, ResNet101V2, MobileNetV2, DenseNet121, DenseNet169, DenseNet201, \
    MobileNetV3Small, ResNet152V2, VGG16, EfficientNetV2B0, EfficientNetV2S, InceptionV3, EfficientNetV2L, \
    EfficientNetV2B3, InceptionResNetV2

from util.model_listutil import ModelListUtil


def type_of_script():
    ipy_str = str(type(get_ipython()))
    if 'zmqshell' in ipy_str:
        return 'jupyter'
    if 'terminal' in ipy_str:
        return 'ipython'

    return 'terminal'


InceptionResNetV2list = [
    # [ModelLayers.ConvBase, InceptionResNetV2, True, 150],
    # [ModelLayers.ConvBase, InceptionResNetV2, True, 200],
    [ModelLayers.ConvBase, InceptionResNetV2, True, 500],
    [ModelLayers.Flatten],
    # [ModelLayers.Dense, 128],
    # [ModelLayers.BatchNormalization],
    [ModelLayers.Dense, 64],
    [ModelLayers.BatchNormalization],
    [ModelLayers.Dense, 32],
    # [ModelLayers.BatchNormalization]
]

MobileNetlist = [
    # [ModelLayers.ConvBase, MobileNetV2, True, 0],
    [ModelLayers.ConvBase, MobileNetV2, True, 40, "avg"],
    # [ModelLayers.ConvBase, MobileNetV2, True, 40],
    [ModelLayers.Flatten],
    # [ModelLayers.Dense, 128],
    # [ModelLayers.BatchNormalization],
    # [ModelLayers.Dense, 64],
    # [ModelLayers.BatchNormalization],
    [ModelLayers.Dense, 32],
]
Inceptionlist = [
    [ModelLayers.ConvBase, InceptionV3, False],
    [ModelLayers.Flatten],
    [ModelLayers.Dense, 128],
    [ModelLayers.BatchNormalization],
    [ModelLayers.Dense, 64],
    [ModelLayers.BatchNormalization],
    [ModelLayers.Dense, 32],
    [ModelLayers.BatchNormalization]
]
DenseNetlist = [
    # [ModelLayers.ConvBase, DenseNet121, True, 150],
    [ModelLayers.ConvBase, DenseNet121, True, 300, "avg"],
    # [ModelLayers.Conv2D, 64, 3, True],
    # [ModelLayers.MaxPooling2D],
    # [ModelLayers.BatchNormalization],
    # [ModelLayers.Flatten],
    [ModelLayers.Dense, 128],
    [ModelLayers.BatchNormalization],
    [ModelLayers.Dense, 64],
    [ModelLayers.BatchNormalization],
    [ModelLayers.Dense, 32],
    # [ModelLayers.BatchNormalization]
]
RESlist = [
    # [ModelLayers.ConvBase, ResNet50V2, True, 300],
    [ModelLayers.ConvBase, ResNet50V2, True, 60],
    # [ModelLayers.BatchNormalization],
    [ModelLayers.Flatten],
    # [ModelLayers.Dense, 128],
    # [ModelLayers.BatchNormalization],
    [ModelLayers.Dense, 64],
    [ModelLayers.BatchNormalization],
    [ModelLayers.Dense, 32],
]

Avglist0 = [
    [ModelLayers.Average],
    [ModelLayers.Dense, 64],
    [ModelLayers.BatchNormalization],
    [ModelLayers.Dense, 32],
    [ModelLayers.BatchNormalization],
]
Avglist = [
    [ModelLayers.Average],
    [ModelLayers.Dense, 64],
    # [ModelLayers.Dense, 64],
    # [ModelLayers.Dense, 64],
    # [ModelLayers.BatchNormalization],
    [ModelLayers.Dense, 32],
    [ModelLayers.BatchNormalization]
]

Avglist1 = [
    [ModelLayers.Average],
    [ModelLayers.Dense, 64],
    [ModelLayers.BatchNormalization],
    [ModelLayers.Dense, 32],
]



def run(model_list: list, conf: ModelConfig, inputs: Input, models: list = None, last_dense: bool = True,
        training: bool = True, batch_size: int = 64):
    m1 = Model(conf, inputs)

    m1.create_by_list_real_cls_functional(model_list, models, last_dense)
    ss = m1.show()
    log.debug(ss.encode('utf-8').decode('gbk'))
    m1.compile()
    if training:
        m1.train(data_manager.x_train, data_manager.y_train, data_manager.x_test, data_manager.y_test,
                 batch_size=batch_size)
    return m1


mobilenetv21 = [
    [ModelLayers.ConvBase, MobileNetV2, False, 0, "avg"],
    [ModelLayers.Dense, 32],
]
mobilenetv2 = [
    [ModelLayers.ConvBase, MobileNetV2, True, 200, "avg"],
    # [ModelLayers.ConvBase, MobileNetV2, True, 0],
    # [ModelLayers.Conv2D, 64, 1],
    # [ModelLayers.ConvBase, MobileNetV2, False, 0, "avg"],
    # [ModelLayers.BatchNormalization],
    # [ModelLayers.Flatten],
    # [ModelLayers.Dense, 128],
    [ModelLayers.Dense, 64],
    [ModelLayers.Dense, 32],
    # [ModelLayers.BatchNormalization],

]

resnet50v2 = [
    # [ModelLayers.ConvBase, ResNet50V2, False, 0, "avg"],
    [ModelLayers.ConvBase, ResNet50V2, True, 1000, "avg"],
    # [ModelLayers.Dense, 64],
    # [ModelLayers.BatchNormalization],
    [ModelLayers.Dense, 32],
    [ModelLayers.BatchNormalization],
]

denenseNet121 = [
    [ModelLayers.ConvBase, DenseNet121, True, 1000, "avg"],
    [ModelLayers.Dense, 32],
]
inceptionv2 = [
    [ModelLayers.ConvBase, InceptionResNetV2, True, 1000, "avg"],
    # [ModelLayers.Dense, 64],
    # [ModelLayers.BatchNormalization],
    [ModelLayers.Dense, 32],
    # [ModelLayers.BatchNormalization],
]


def main():
    print(log.handlers)
    conf = ModelConfig()
    log.debug("config:" + str(conf))
    log.info("running at " + type_of_script())
    conf.LEARNING_RATE = 0.001/5
    conf.EPOCHS = 20  # 测试
    conf.OPTIMIZER = SGD
    # conf.EPOCHS = 20
    # conf.DECAY = 0.005
    # conf.DECAY = 0.0025
    conf.DECAY = 0.01

    # mtest = run(mobilenetv2, conf, inputs)
    # conf.DECAY = 0.005
    # conf.LEARNING_RATE=0.0005
    conf.IMAGE_HEIGHT = Config.IMAGE_HEIGHT
    conf.IMAGE_WIDTH = Config.IMAGE_WIDTH
    inputs = Input(shape=(conf.IMAGE_HEIGHT, conf.IMAGE_HEIGHT, conf.COLOR_CHANNELS))
    # m1 = run(MobileNetlist, conf, inputs)
    # # m1 = run(InceptionResNetV2list, conf, inputs)
    # m11 = keras.Model(inputs, m1.model.layers[-1].output)
    # conf.EPOCHS = 20
    # m2 = run(RESlist, conf, inputs)
    # m21 = keras.Model(inputs, m2.model.layers[-1].output)
    # conf.DECAY = 0.01
    # m3 = run(DenseNetlist, conf, inputs)
    # m31 = keras.Model(inputs, m3.model.layers[-1].output)
    # # # # # m3.evaluation(data_manager.x_test, data_manager.y_test)
    # conf.EPOCHS = 4
    # m4 = run(InceptionResNetV2list, conf, inputs)
    # mtest = run(inceptionv2, conf, inputs)
    mtest = run(mobilenetv2, conf, inputs)
    # mtest = run(denenseNet121, conf, inputs)
    # mtest = run(resnet50v2, conf, inputs)


if __name__ == "__main__":
    main()
    # conf = ModelConfig()
    # conf.IMAGE_HEIGHT = Config.IMAGE_HEIGHT
    # conf.IMAGE_WIDTH = Config.IMAGE_WIDTH
    # conf.LEARNING_RATE = 0.001 / 5
    # conf.DECAY = 0.01
    #
    # inputs = Input(shape=(conf.IMAGE_HEIGHT, conf.IMAGE_HEIGHT, conf.COLOR_CHANNELS))
    #
    # # # m1 = Model(conf, inputs)  # 0.9962
    # # # model_list = ModelListUtil.get_model_list_from_db(521)
    # # # m1.load(model_list, f"storage/{521}.h5")
    # # #
    # # # m2 = Model(conf, inputs)  # 0.9766
    # # # model_list = ModelListUtil.get_model_list_from_db(539)
    # # # m2.load(model_list, f"storage/{539}.h5")
    # # #
    # # # # m3 = Model(conf, inputs)#0.9922
    # # # # model_list = ModelListUtil.get_model_list_from_db(519)
    # # # # m3.load(model_list, f"storage/{519}.h5")
    # # # m3 = Model(conf, inputs)  # 0.9922
    # # # model_list = ModelListUtil.get_model_list_from_db(576)
    # # # m3.load(model_list, f"storage/{576}.h5")
    # # # m4 = Model(conf, inputs)  # 0.9796
    # # # model_list = ModelListUtil.get_model_list_from_db(523)
    # # # m4.load(model_list, f"storage/{523}.h5")
    # #
    # m1 = Model(conf, inputs)
    # model_list = ModelListUtil.get_model_list_from_db(620)
    # m1.load(model_list, f"storage/{620}.h5")
    #
    # m2 = Model(conf, inputs)#0.9766
    # model_list = ModelListUtil.get_model_list_from_db(502)
    # m2.load(model_list, f"storage/{502}.h5")
    #
    # m3 = Model(conf, inputs)#0.9922
    # model_list = ModelListUtil.get_model_list_from_db(504)
    # m3.load(model_list, f"storage/{504}.h5")
    #
    # m4 = Model(conf, inputs)#0.9796
    # model_list = ModelListUtil.get_model_list_from_db(505)
    # m4.load(model_list, f"storage/{505}.h5")
    #
    # # log.info("m1")
    # # m1.evaluation(data_manager.x_test, data_manager.y_test)  # 0.9855
    # # # log.info("m2")
    # # # m2.evaluation(data_manager.x_test, data_manager.y_test)  # 0.9809
    # # # log.info("m3")
    # # # m3.evaluation(data_manager.x_test, data_manager.y_test)  # 0.9769
    # # # log.info("m4")
    # # # m4.evaluation(data_manager.x_test, data_manager.y_test)  # 0.9894
    # #
    # m11 = keras.Model(inputs, m1.model.layers[-1].output)
    # m21 = keras.Model(inputs, m2.model.layers[-1].output)
    # m31 = keras.Model(inputs, m3.model.layers[-1].output)
    # m41 = keras.Model(inputs, m4.model.layers[-1].output)
    # # # m11 = keras.Model(inputs, m1.model.layers[-1].output)
    # # # m21 = keras.Model(inputs, m2.model.layers[-1].output)
    # # # m31 = keras.Model(inputs, m3.model.layers[-1].output)
    # # # m41 = keras.Model(inputs, m4.model.layers[-1].output)
    # # # m11.layers[1].trainable = False
    # # # m21.layers[1].trainable = False
    # # # m31.layers[1].trainable = False
    # # # m41.layers[1].trainable = False
    # # # m1.trainable = True
    # # # m2.trainable = True
    # # # m3.trainable = True
    # # # m4.trainable = True
    # # # m1.trainable = False
    # # # m2.trainable = False
    # # # m3.trainable = False
    # # # m4.trainable = False
    # m11.trainable = False
    # m21.trainable = False
    # m31.trainable = False
    # m41.trainable = False
    # # # for i in range(1):
    # # #     m11.layers[-1-i].trainable = True
    # # #     m21.layers[-1-i].trainable = True
    # # #     m31.layers[-1-i].trainable = True
    # # #     m41.layers[-1-i].trainable = True
    # # # conf.DECAY = 0.01
    # conf.EPOCHS = 20
    # # # conf.LEARNING_RATE /= 10
    # # # # m = run(Ensemblelist, conf, inputs, [m11, m21, m31, m41])
    # m = run(Avglist, conf, inputs, [m11, m21, m31, m41])
    # # with open("MobileNet50V2结构图.txt","w") as f:
    # #     f.write(m1.show().encode("utf-8").decode("gbk"))
    # #
    # # plot_model(m1.model.layers[1],
    # #            show_shapes=True,
    # #            to_file="shape.png",
    # #            show_layer_names=True,
    # #            show_layer_activations=True)
    # # with open("MobileNet50V2 state dict.txt","w") as f:
    # #     f.write(str(m1.model.state_dict()))
    #
    # # m = run(Avglist, conf, inputs, [m11, m21, m31, m41], last_dense=False, training=False, batch_size=128)
    # # m = run(Avglist, conf, inputs, [m11, m21, m31, m41], last_dense=False, training=False)
    # # # m = run(Avglist, conf, inputs, [m11, m21, m31])
    # # # m = run(Maxlist, conf, inputs, [m11, m21, m31, m41])
    # m.evaluation(data_manager.x_test, data_manager.y_test,128)
    # #
    # #
    # #
    # #
    # # # t = TrainData()
    # # # t.decay = conf.DECAY
    # # # t.learning_rate = conf.LEARNING_RATE
    # # # t.image_width = conf.IMAGE_WIDTH
    # # # t.image_height = conf.IMAGE_HEIGHT
    # # # t.count_epoch = -1
    # # # td = TrainDataDao()
    # # #
    # # # idx = td.setExcludeImgName(t, "not exist")
    # # #
    # # # p = Predict()
    # # # p.calc_metrics(m.model, data_manager.x_test, data_manager.y_test, -1, idx)
    # #
    # # # conf = ModelConfig()
    # # # inputs = Input(shape=(conf.IMAGE_HEIGHT, conf.IMAGE_HEIGHT, conf.COLOR_CHANNELS))
    # # #
    # # # m1= Model(conf, inputs)
    # # # m1.load(mobilenetv2,"storage/332.h5")
    # # # m1.evaluation(data_manager.x_test,data_manager.y_test)
